import torch
from train.metrics import MCR_MSE

"""
对模型训练和验证
调用一次即一个epoch训练或者验证
"""


def train_network(model, optimizer, device, data_loader, config):
    model.train()
    iter = 0
    epoch_loss = 0  # 计算本次epoch损失
    epoch_assess = 0  # 计算本次epoch评价指标
    for iter, (batch_graphs, batch_labels) in enumerate(data_loader):
        batch_graphs = batch_graphs.to(device)
        batch_mask = batch_graphs.ndata['mask'].to(device)  # Mask信息(指定需要训练的部分)
        batch_x = batch_graphs.ndata['nfeat'].to(device)  # 节点特征
        batch_e = batch_graphs.edata['efeat'].to(device)  # 边特征
        batch_labels = batch_labels.to(device)  # 真实标签

        optimizer.zero_grad()
        if config.add_loop_type == True:  # 若使用loop_type特征
            batch_loop = batch_graphs.ndata['loop_type'].to(device)
            if config.pos_enc == True:  # 若使用位置编码
                batch_pos_enc = batch_graphs.ndata['pos_enc'].to(device)
                # sign_flip = torch.rand(batch_pos_enc.size(1)).to(device)
                # sign_flip[sign_flip >= 0.5] = 1.0
                # sign_flip[sign_flip < 0.5] = -1.0
                # batch_pos_enc = batch_pos_enc * sign_flip.unsqueeze(0)
                batch_scores = model.forward(batch_graphs, batch_x, batch_e, loop_type=batch_loop,
                                             h_pos_enc=batch_pos_enc)
            else:
                batch_scores = model.forward(batch_graphs, batch_x, batch_e, loop_type=batch_loop)
        else:
            if config.pos_enc == True:  # 若使用位置编码
                batch_pos_enc = batch_graphs.ndata['pos_enc'].to(device)
                # sign_flip = torch.rand(batch_pos_enc.size(1)).to(device)
                # sign_flip[sign_flip >= 0.5] = 1.0
                # sign_flip[sign_flip < 0.5] = -1.0
                # batch_pos_enc = batch_pos_enc * sign_flip.unsqueeze(0)
                batch_scores = model.forward(batch_graphs, batch_x, batch_e, h_pos_enc=batch_pos_enc)
            else:
                batch_scores = model.forward(batch_graphs, batch_x, batch_e, )

        loss = model.loss(batch_scores, batch_labels, batch_mask)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.detach().item()
        epoch_assess += MCR_MSE(batch_scores, batch_labels, batch_mask)
    epoch_loss /= (iter + 1)
    epoch_assess /= (iter + 1)
    return epoch_loss, epoch_assess, optimizer


def evaluate_network(model, device, data_loader, config):
    model.eval()
    epoch_test_loss = 0  # 计算本次epoch损失
    epoch_test_assess = 0  # 计算本次epoch评价指标
    with torch.no_grad():
        for iter, (batch_graphs, batch_labels) in enumerate(data_loader):
            batch_graphs = batch_graphs.to(device)
            batch_mask = batch_graphs.ndata['mask'].to(device)  # Mask信息(指定需要训练的部分)
            batch_x = batch_graphs.ndata['nfeat'].to(device)  # 节点特征
            batch_e = batch_graphs.edata['efeat'].to(device)  # 边特征
            batch_labels = batch_labels.to(device)  # 真实标签

            if config.add_loop_type == True:  # 若使用loop_type特征
                batch_loop = batch_graphs.ndata['loop_type'].to(device)
                if config.pos_enc == True:  # 若使用位置编码
                    batch_pos_enc = batch_graphs.ndata['pos_enc'].to(device)
                    # sign_flip = torch.rand(batch_pos_enc.size(1)).to(device)
                    # sign_flip[sign_flip >= 0.5] = 1.0
                    # sign_flip[sign_flip < 0.5] = -1.0
                    # batch_pos_enc = batch_pos_enc * sign_flip.unsqueeze(0)
                    batch_scores = model.forward(batch_graphs, batch_x, batch_e, loop_type=batch_loop,
                                                 h_pos_enc=batch_pos_enc)
                else:
                    batch_scores = model.forward(batch_graphs, batch_x, batch_e, loop_type=batch_loop)
            else:
                if config.pos_enc == True:  # 若使用位置编码
                    batch_pos_enc = batch_graphs.ndata['pos_enc'].to(device)
                    # sign_flip = torch.rand(batch_pos_enc.size(1)).to(device)
                    # sign_flip[sign_flip >= 0.5] = 1.0
                    # sign_flip[sign_flip < 0.5] = -1.0
                    # batch_pos_enc = batch_pos_enc * sign_flip.unsqueeze(0)
                    batch_scores = model.forward(batch_graphs, batch_x, batch_e, h_pos_enc=batch_pos_enc)
                else:
                    batch_scores = model.forward(batch_graphs, batch_x, batch_e)

            loss = model.loss(batch_scores, batch_labels, batch_mask)
            epoch_test_loss += loss.detach().item()
            epoch_test_assess += MCR_MSE(batch_scores, batch_labels, batch_mask)
        epoch_test_loss /= (iter + 1)
        epoch_test_assess /= (iter + 1)
    return epoch_test_loss, epoch_test_assess
